"""
Spherical coordinate grid implementations.
"""

using LinearAlgebra
using StaticArrays

"""
    PolarSymGrid <: AbstractSphericalGrid

2D polar coordinate grid with angular symmetry (r-only).

# Fields
- `num_r::Int`: Number of radial points
- `r_bounds::Tuple{Float64,Float64}`: Radial bounds [r_min, r_max]

# Examples
```julia
# Polar grid from r=0 to r=5 with 64 radial points
grid = PolarSymGrid(64, (0.0, 5.0))
```
"""
struct PolarSymGrid <: AbstractSphericalGrid
    num_r::Int
    r_bounds::Tuple{Float64,Float64}
    
    function PolarSymGrid(num_r::Int, r_bounds::Tuple{Float64,Float64})
        num_r > 0 || throw(ArgumentError("Number of radial points must be positive"))
        r_bounds[1] >= 0 || throw(ArgumentError("Minimum radius must be non-negative"))
        r_bounds[1] < r_bounds[2] || throw(ArgumentError("Invalid radial bounds"))
        new(num_r, r_bounds)
    end
end

# Interface implementation for PolarSymGrid
ndim(::PolarSymGrid) = 1  # Effective dimension due to symmetry
Base.size(grid::PolarSymGrid) = (grid.num_r,)
extent(grid::PolarSymGrid) = (grid.r_bounds,)

function spacing(grid::PolarSymGrid)
    dr = (grid.r_bounds[2] - grid.r_bounds[1]) / (grid.num_r - 1)
    (dr,)
end

function coordinate(grid::PolarSymGrid, idx)
    r = grid.r_bounds[1] + (idx[1] - 1) * spacing(grid)[1]
    SVector(r)
end

function index(grid::PolarSymGrid, coord)
    r = coord[1]
    dr = spacing(grid)[1]
    idx_r = round(Int, (r - grid.r_bounds[1]) / dr + 1)
    (idx_r,)
end

"""
    SphericalSymGrid <: AbstractSphericalGrid

3D spherical coordinate grid with spherical symmetry (r-only).

# Fields
- `num_r::Int`: Number of radial points
- `r_bounds::Tuple{Float64,Float64}`: Radial bounds [r_min, r_max]

# Examples
```julia
# Spherical grid from r=0.1 to r=10 with 128 radial points
grid = SphericalSymGrid(128, (0.1, 10.0))
```
"""
struct SphericalSymGrid <: AbstractSphericalGrid
    num_r::Int
    r_bounds::Tuple{Float64,Float64}
    
    function SphericalSymGrid(num_r::Int, r_bounds::Tuple{Float64,Float64})
        num_r > 0 || throw(ArgumentError("Number of radial points must be positive"))
        r_bounds[1] >= 0 || throw(ArgumentError("Minimum radius must be non-negative"))
        r_bounds[1] < r_bounds[2] || throw(ArgumentError("Invalid radial bounds"))
        new(num_r, r_bounds)
    end
end

# Interface implementation for SphericalSymGrid
ndim(::SphericalSymGrid) = 1  # Effective dimension due to symmetry
Base.size(grid::SphericalSymGrid) = (grid.num_r,)
extent(grid::SphericalSymGrid) = (grid.r_bounds,)

function spacing(grid::SphericalSymGrid)
    dr = (grid.r_bounds[2] - grid.r_bounds[1]) / (grid.num_r - 1)
    (dr,)
end

function coordinate(grid::SphericalSymGrid, idx)
    r = grid.r_bounds[1] + (idx[1] - 1) * spacing(grid)[1]
    SVector(r)
end

function index(grid::SphericalSymGrid, coord)
    r = coord[1]
    dr = spacing(grid)[1]
    idx_r = round(Int, (r - grid.r_bounds[1]) / dr + 1)
    (idx_r,)
end

# Coordinate transformation utilities
"""
    spherical_to_cartesian(r, θ, φ)

Convert spherical coordinates (r, θ, φ) to Cartesian (x, y, z).
"""
function spherical_to_cartesian(r, θ, φ)
    x = r * sin(θ) * cos(φ)
    y = r * sin(θ) * sin(φ)
    z = r * cos(θ)
    return SVector(x, y, z)
end

"""
    cartesian_to_spherical(x, y, z)

Convert Cartesian coordinates (x, y, z) to spherical (r, θ, φ).
"""
function cartesian_to_spherical(x, y, z)
    r = sqrt(x^2 + y^2 + z^2)
    θ = acos(z / r)
    φ = atan(y, x)
    return SVector(r, θ, φ)
end

"""
    polar_to_cartesian(r, φ)

Convert polar coordinates (r, φ) to Cartesian (x, y).
"""
function polar_to_cartesian(r, φ)
    x = r * cos(φ)
    y = r * sin(φ)
    return SVector(x, y)
end

"""
    cartesian_to_polar(x, y)

Convert Cartesian coordinates (x, y) to polar (r, φ).
"""
function cartesian_to_polar(x, y)
    r = sqrt(x^2 + y^2)
    φ = atan(y, x)
    return SVector(r, φ)
end

# Volume elements for integration
"""
    volume_element(grid::PolarSymGrid, r::Float64)

Volume element for polar coordinates: 2π * r
"""
function volume_element(grid::PolarSymGrid, r::Float64)
    return 2π * r
end

"""
    volume_element(grid::SphericalSymGrid, r::Float64)

Volume element for spherical coordinates: 4π * r²
"""
function volume_element(grid::SphericalSymGrid, r::Float64)
    return 4π * r^2
end

# Integration support
"""
    integrate(grid::PolarSymGrid, values::AbstractVector)

Integrate values over polar grid with proper volume weighting.
"""
function integrate(grid::PolarSymGrid, values::AbstractVector)
    rs = [coordinate(grid, (i,))[1] for i in 1:grid.num_r]
    weights = [volume_element(grid, r) for r in rs]
    dr = spacing(grid)[1]
    return sum(values .* weights) * dr
end

"""
    integrate(grid::SphericalSymGrid, values::AbstractVector)

Integrate values over spherical grid with proper volume weighting.
"""
function integrate(grid::SphericalSymGrid, values::AbstractVector)
    rs = [coordinate(grid, (i,))[1] for i in 1:grid.num_r]
    weights = [volume_element(grid, r) for r in rs]
    dr = spacing(grid)[1]
    return sum(values .* weights) * dr
end